﻿import pygame
import sys
import gymnasium as gym
from gymnasium import spaces
import numpy as np
from Game.G3Maze.record import GameRecorder

# 颜色常量保持不变
COLORS = {
    "background": (255, 255, 255),
    "grid": (200, 200, 200),
    "player": (255, 0, 0),
    "wall": (0, 0, 255),
    "exit": (0, 255, 0),
    "reward": (255, 215, 0),
}


class MazeEnv(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 60}

    def __init__(self, render_mode=None):
        # 游戏常量设置
        self.TILE_SIZE = 100
        self.MAP_ROWS, self.MAP_COLS = 3, 3
        self.WINDOW_SIZE = (
            self.MAP_COLS * self.TILE_SIZE,
            self.MAP_ROWS * self.TILE_SIZE,
        )

        # 强化学习接口
        self.observation_space = spaces.Dict(
            {
                "player": spaces.Box(0, 2, shape=(2,), dtype=int),
                "rewards": spaces.MultiBinary((self.MAP_ROWS, self.MAP_COLS)),
            }
        )
        self.action_space = spaces.Discrete(4)  # 0:左,1:右,2:上,3:下
        self.render_mode = render_mode

        # 图形系统初始化
        self.window = None
        self.clock = None
        self.font = None

        # 游戏逻辑初始化
        self._init_game_params()

        self.recorder = None  # 将由外部注入

    def set_recorder(self, recorder):
        """依赖注入记录器"""
        self.recorder = recorder

    def _init_game_params(self):
        """初始化游戏参数"""
        self.exits = {(2, 0, "left"), (2, 2, "right")}
        self.walls = {
            (0, 1, "right"),
            (1, 0, "down"),
            (1, 2, "down"),
        }
        self.original_rewards = {(2, 0): 10, (2, 2): 20}
        self.player_pos = np.array([0, 1], dtype=int)
        self.active_rewards = set()
        self.score = 0

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)

        # 重置游戏状态
        self.player_pos = np.array([0, 1], dtype=int)
        self.active_rewards = set(self.original_rewards.keys())
        self.score = 0

        # 初始化渲染系统
        if self.render_mode == "human":
            self._init_render()

        return self._get_obs(), {"score": self.score}

    def step(self, action):
        direction_map = [
            ("left", 0, -1, "target", "right"),
            ("right", 0, 1, "current", "right"),
            ("up", -1, 0, "target", "down"),
            ("down", 1, 0, "current", "down"),
        ]
        dir_name, dr, dc, check_pos, wall_dir = direction_map[action]

        # 1. 获取移动前状态（必须最先执行）
        current_r, current_c = self.player_pos
        current_state = (int(current_r), int(current_c))
    
        # 2. 计算理论上的新位置
        new_r = current_r + dr
        new_c = current_c + dc
        theoretical_next_state = (int(new_r), int(new_c))

        # 3. 检查移动有效性
        valid = self._is_valid_move(new_r, new_c, check_pos, wall_dir)
        terminated = (current_r, current_c, dir_name) in self.exits
        reward = 0
        actual_next_state = current_state  # 默认无效移动时保持原位

        # 4. 处理有效移动
        if valid and not terminated:
            actual_next_state = theoretical_next_state
            self.player_pos = np.array([new_r, new_c], dtype=int)

            # 检查奖励
            if (new_r, new_c) in self.active_rewards:
                reward = self.original_rewards[(new_r, new_c)]
                self.active_rewards.remove((new_r, new_c))
                self.score += reward

        # 5. 记录状态转移（在所有状态变更之后）
        if self.recorder:
            self.recorder.record(
                state=current_state,
                action=int(action),
                next_state=actual_next_state,
                reward=float(reward),
                done=bool(terminated),
            )

        # 6. 渲染处理
        if self.render_mode == "human":
            self.render()

        return self._get_obs(), reward, terminated, False, {"score": self.score}

    def _get_obs(self):
        return {
            "player": self.player_pos.copy(),
            "rewards": self._get_reward_grid(),
        }

    def _get_reward_grid(self):
        grid = np.zeros((self.MAP_ROWS, self.MAP_COLS), dtype=int)
        for r, c in self.active_rewards:
            grid[r, c] = 1
        return grid

    def _is_valid_move(self, new_r, new_c, check_pos, wall_dir):
        """验证移动逻辑"""
        if not (0 <= new_r < self.MAP_ROWS and 0 <= new_c < self.MAP_COLS):
            return False

        # 检查墙壁
        if check_pos == "target":
            check_r, check_c = new_r, new_c
        else:
            check_r, check_c = self.player_pos

        wall_exists = (check_r, check_c, wall_dir) in self.walls
        return not wall_exists

    # 图形渲染核心 ------------------------------------------------
    def _init_render(self):
        """初始化图形系统"""
        if self.window is None:
            pygame.init()
            self.window = pygame.display.set_mode(self.WINDOW_SIZE)
            pygame.display.set_caption("MazeRL - Gymnasium")
            self.font = pygame.font.Font(None, 36)
            self.clock = pygame.time.Clock()

    def render(self):
        """完整图形渲染逻辑"""
        if self.render_mode is None:
            return

        # 初始化检查
        if self.window is None and self.render_mode == "human":
            self._init_render()

        # 处理退出事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.close()
                sys.exit()

        # 绘制游戏元素
        self.window.fill(COLORS["background"])
        self._draw_grid()
        self._draw_exits()
        self._draw_walls()
        self._draw_rewards()
        self._draw_player()
        self._draw_score()

        # 更新显示
        if self.render_mode == "human":
            pygame.display.flip()
            self.clock.tick(self.metadata["render_fps"])
        elif self.render_mode == "rgb_array":
            return np.transpose(
                np.array(pygame.surfarray.pixels3d(self.window)), axes=(1, 0, 2)
            )

    def _draw_grid(self):
        """绘制网格线"""
        # 垂直网格线
        for x in range(self.TILE_SIZE, self.WINDOW_SIZE[0], self.TILE_SIZE):
            pygame.draw.line(
                self.window, COLORS["grid"], (x, 0), (x, self.WINDOW_SIZE[1]), 2
            )
        # 水平网格线
        for y in range(self.TILE_SIZE, self.WINDOW_SIZE[1], self.TILE_SIZE):
            pygame.draw.line(
                self.window, COLORS["grid"], (0, y), (self.WINDOW_SIZE[0], y), 2
            )

    def _draw_walls(self):
        """绘制所有墙壁"""
        for r, c, d in self.walls:
            if d == "right":
                x = (c + 1) * self.TILE_SIZE
                pygame.draw.line(
                    self.window,
                    COLORS["wall"],
                    (x, r * self.TILE_SIZE),
                    (x, (r + 1) * self.TILE_SIZE),
                    8,
                )
            elif d == "down":
                y = (r + 1) * self.TILE_SIZE
                pygame.draw.line(
                    self.window,
                    COLORS["wall"],
                    (c * self.TILE_SIZE, y),
                    ((c + 1) * self.TILE_SIZE, y),
                    8,
                )

    def _draw_exits(self):
        """绘制出口"""
        for r, c, d in self.exits:
            x = c * self.TILE_SIZE + (self.TILE_SIZE if d == "right" else 0)
            pygame.draw.line(
                self.window,
                COLORS["exit"],
                (x, r * self.TILE_SIZE),
                (x, (r + 1) * self.TILE_SIZE),
                8,
            )

    def _draw_player(self):
        """绘制玩家角色"""
        r, c = self.player_pos
        center = (
            c * self.TILE_SIZE + self.TILE_SIZE // 2,
            r * self.TILE_SIZE + self.TILE_SIZE // 2,
        )
        pygame.draw.circle(self.window, COLORS["player"], center, self.TILE_SIZE // 3)

    def _draw_rewards(self):
        """绘制奖励物品"""
        for r, c in self.active_rewards:
            center = (
                c * self.TILE_SIZE + self.TILE_SIZE // 2,
                r * self.TILE_SIZE + self.TILE_SIZE // 2,
            )
            pygame.draw.rect(
                self.window,
                COLORS["reward"],
                (center[0] - 20, center[1] - 20, 40, 40),
                border_radius=8,
            )
            text = self.font.render(str(self.original_rewards[(r, c)]), True, (0, 0, 0))
            self.window.blit(text, (center[0] - 10, center[1] - 10))

    def _draw_score(self):
        """绘制分数显示"""
        score_text = self.font.render(f"Score: {self.score}", True, (0, 0, 0))
        self.window.blit(score_text, (10, 10))

    def close(self):
        """关闭环境"""
        if self.window is not None:
            pygame.quit()
            self.window = None


# 使用示例
if __name__ == "__main__":
    env = MazeEnv(render_mode="human")
    # 创建记录器并注入环境
    recorder = GameRecorder()
    env.set_recorder(recorder)

    for episode in range(3):  # 测试3局游戏
        print(f"\n=== 开始第 {episode+1} 局游戏 ===")

        obs, info = env.reset()
        recorder.new_episode()  # 明确开始新记录

        terminated = False
        running = True
        while running and not terminated:
            action = None
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_LEFT:
                        action = 0
                    elif event.key == pygame.K_RIGHT:
                        action = 1
                    elif event.key == pygame.K_UP:
                        action = 2
                    elif event.key == pygame.K_DOWN:
                        action = 3
                    elif event.key == pygame.K_ESCAPE:  # 退出快捷键
                        running = False

            if action is not None:
                obs, reward, terminated, truncated, info = env.step(action)
                # 显示当前状态
                print(
                    f"动作: {['左','右','上','下'][action]} | "
                    f"位置: {tuple(env.player_pos.tolist())} | "
                    f"奖励: {reward} | "
                    f"得分: {info['score']}"
                )

                if terminated:
                    print(f"=== 第 {episode+1} 局结束！得分: {info['score']} ===")
            env.render()
            pygame.time.delay(100)  # 添加延迟使游戏可玩

    env.close()
